﻿#pragma once
#include <iostream>
using namespace std;

// 枚举值表示颜色
enum Colour
{
    RED,
    BLACK
};

// 节点结构
template<class T>
struct RBTreeNode
{
    // 这里更新控制平衡也需要parent指针
    T _data;
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    Colour _col;        //颜色 红黑树的关键

    RBTreeNode(const T& data)
        :_data(data)
        , _left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
    {
    }
};

// 迭代器
template<class T, class Ref, class Ptr>
struct TreeIterator
{
    typedef RBTreeNode<T> Node;
    typedef TreeIterator<T, Ref, Ptr> Self;
    typedef TreeIterator<T, T&, T*> Iterator;

    Node* _node;

    TreeIterator(Node* node)
        :_node(node)
    {
    }

    // 允许从普通迭代器构造const迭代器
    TreeIterator(const Iterator& it)
        :_node(it._node)
    {
    }

    Ref operator* () const
    {
        return _node->_data;
    }

    Ptr operator->() const
    {
        return &_node->_data;
    }

    bool operator!=(const Self& s) const
    {
        return _node != s._node;
    }

    bool operator==(const Self& s) const
    {
        return _node == s._node;
    }

    // 前置--
    Self& operator--()
    {
        if (_node->_left)
        {
            // 左子树存在，找左子树的最右节点
            Node* max = _node->_left;
            while (max->_right)
            {
                max = max->_right;
            }
            _node = max;
        }
        else
        {
            // 左子树不存在，向上找第一个是父节点右孩子的节点
            Node* cur = _node;
            Node* parent = _node->_parent;
            while (parent && cur == parent->_left)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    // 后置--
    Self operator--(int)
    {
        Self tmp = *this;
        --(*this);
        return tmp;
    }

    // 前置++
    Self& operator++()
    {
        // 当前右不为空，下一个就是右子树有序第一个（最左结点）
        if (_node->_right)
        {
            Node* min = _node->_right;
            while (min->_left)
            {
                min = min->_left;
            }
            _node = min;
        }
        // 当前右为空，下一个就是孩子是父亲左的那个祖先结点
        else
        {
            Node* cur = _node;
            Node* parent = _node->_parent;
            while (parent && cur == parent->_right)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    // 后置++
    Self operator++(int)
    {
        Self tmp = *this;
        ++(*this);
        return tmp;
    }
};

template<class K, class T, class KeyOfT>
class RBTree
{
    typedef RBTreeNode<T> Node;
public:
    typedef TreeIterator<T, T&, T*> Iterator;
    typedef TreeIterator<T, const T&, const T*> ConstIterator;

    Iterator Begin()
    {
        Node* min = _root;
        while (min && min->_left)
        {
            min = min->_left;
        }
        return Iterator(min);
    }

    Iterator End()
    {
        return Iterator(nullptr);
    }

    ConstIterator Begin() const
    {
        Node* min = _root;
        while (min && min->_left)
        {
            min = min->_left;
        }
        return ConstIterator(min);
    }

    ConstIterator End() const
    {
        return ConstIterator(nullptr);
    }

    // 查找节点并返回const迭代器
    ConstIterator Find(const K& key) const
    {
        KeyOfT kot;
        Node* cur = _root;
        while (cur)
        {
            if (kot(cur->_data) < key)
            {
                cur = cur->_right;
            }
            else if (kot(cur->_data) > key)
            {
                cur = cur->_left;
            }
            else
            {
                return ConstIterator(cur);
            }
        }
        return End();
    }

    // 红黑树插入
    // 旋转代码的实现跟AVL树是一样的，只是不需要更新平衡因子
    pair<Iterator, bool> Insert(const T& data)
    {
        if (_root == nullptr)
        {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(Iterator(_root), true);
        }

        KeyOfT kot;
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur)
        {
            if (kot(cur->_data) < kot(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (kot(cur->_data) > kot(data))
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return make_pair(Iterator(cur), false);
            }
        }

        cur = new Node(data);
        // 新增节点。颜色红色给红色
        cur->_col = RED;
        if (kot(parent->_data) < kot(data))
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }
        cur->_parent = parent;

        while (parent && parent->_col == RED)
        {
            Node* grandfather = parent->_parent;
            // g
            // p u
            if (parent == grandfather->_left)
            {
                Node* uncle = grandfather->_right;
                if (uncle && uncle->_col == RED)
                {
                    // 叔叔存在且为红 -》变色
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    //继续往上处理
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else
                {
                    // 叔叔存在且为黑或不存在-》旋转+变色
                    if (cur == parent->_left)
                    {
                        // g
                        // p u
                        //c
                        //单旋
                        RotateR(grandfather);
                        //变色
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else
                    {
                        // g
                        // p u
                        // c
                        //双旋
                        RotateL(parent);
                        RotateR(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
            else
            {
                // g
                // u p
                Node* uncle = grandfather->_left;
                // 叔叔存在且为红色-》变色
                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    // 继续往上处理
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else // 情况⼆：叔叔不存在或者存在为黑
                {
                    // 情况⼆：叔叔不存在或者存在为黑
                    // 旋转+变色
                    // g
                    // u p
                    // c
                    if (cur == parent->_right)
                    {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else
                    {
                        // g
                        // u p
                        // c
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return make_pair(Iterator(cur), true);
    }

private:
    // 右单旋
    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;

        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;

        Node* parentParent = parent->_parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            if (parentParent->_left == parent)
            {
                parentParent->_left = subL;
            }
            else
            {
                parentParent->_right = subL;
            }

            subL->_parent = parentParent;
        }
    }

    // 左单旋
    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;

        parent->_right = subRL;
        if (subRL)
            subRL->_parent = parent;

        Node* parentParent = parent->_parent;

        subR->_left = parent;
        parent->_parent = subR;

        if (parent == _root)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            if (parentParent->_left == parent)
            {
                parentParent->_left = subR;
            }
            else
            {
                parentParent->_right = subR;
            }

            subR->_parent = parentParent;
        }
    }

    Node* _root = nullptr;
};